import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation


# 定义函数
def f(x):
    return x ** 2


# 导数
def df(x):
    return 2 * x


# 梯度下降参数
lr = 0.1
# 迭代次数
n_iterations = 10
# 初始值
x1 = 2.5

# 梯度下降算法
# for i in range(n_iterations):
#     gradient = df(x1)
#     x1 = x1 - lr * gradient

# 绘制原始函数
x = np.linspace(-3, 3, 100)

y = f(x)
# plt.figure(figsize=(8, 6))
# plt.plot(x, y, label='f(x)=x^2')


fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(x, y, label='f(x)=x^2')

point, = ax.plot([], [], 'ro', label='Gradient Descent')
value_display = ax.text(0.7, 0.02, '', transform=ax.transAxes)


# # 绘制梯度下降过程中x 的位置
# x_history = []
# y_history = []


def init():
    point.set_data([], [])
    value_display.set_text('')
    return point, value_display


def update(i):
    global x1
    gradient = df(x1)
    x1 -= lr * gradient
    point.set_data(x,f(x1))
    value_display.set_text('Min={:.2f},{:.2f}'.format(x1, f(x1)))
    return point, value_display


ani = FuncAnimation(fig, update, frames=np.arange(0, n_iterations), init_func=init, blit=True)

ax.legend()
ax.set_xlabel('x')
ax.set_ylabel('f(x)')
ax.set_title('Function and Gradient Descent  Animation')
ax.grid(True)

plt.show()

